package com.zjw.nanyinge.service.impl;

import com.theokanning.openai.completion.CompletionChoice;
import com.theokanning.openai.completion.CompletionRequest;
import com.theokanning.openai.completion.chat.ChatCompletionChoice;
import com.theokanning.openai.completion.chat.ChatCompletionRequest;
import com.theokanning.openai.completion.chat.ChatCompletionResult;
import com.theokanning.openai.completion.chat.ChatMessage;
import com.theokanning.openai.edit.EditChoice;
import com.theokanning.openai.edit.EditRequest;
import com.theokanning.openai.edit.EditResult;
import com.theokanning.openai.image.CreateImageRequest;
import com.theokanning.openai.image.Image;
import com.theokanning.openai.service.OpenAiService;
import com.zjw.nanyinge.service.OpenAIChatService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

@Slf4j
@Service
public class OpenAIChatServiceImpl implements OpenAIChatService {
    @Resource
    private OpenAiService openAiService;

    @Value("${openAI.completion.model.three}")
    private String completionModelThree;
    @Value("${openAI.completion.model.threePlus}")
    private String completionModelThreePlus;
    @Value("${openAI.completion.model.four}")
    private String completionModelFour;
    @Value("${openAI.edit.model}")
    private String editModel;


    public String createCompletion(String text) {
        if (!text.startsWith("Human")) {
            text = "Human:" + text + "\n\nAI:";
        }
        if (text.length() > 150) {
            text = text.substring(text.length() - 100, text.length() - 1);
        }
        CompletionRequest completionRequest = CompletionRequest.builder()
                .user("zhaojinwei")
                .prompt(text)
                .model(completionModelThree)
                .temperature(0.1d)
                .maxTokens(120)
                .stop(Arrays.asList("Human:", "AI:"))
                .echo(true)
                .n(1)
                .build();
        List<CompletionChoice> choices = openAiService.createCompletion(completionRequest).getChoices();
        for (CompletionChoice choice : choices) {
            return choice.getText();
        }
        return "";
    }

    public String createChatCompletion(String content) {
        log.info("openAI模型：{}", completionModelThreePlus);
        if (!content.startsWith("Human")) {
            content = "Human:" + content + "\n\nAI:";
        }
        List<ChatMessage> messages = new ArrayList<>();
        ChatMessage chatMessage = new ChatMessage();
        chatMessage.setContent(content);
        chatMessage.setRole("user");
        messages.add(chatMessage);
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(completionModelThreePlus)
                .messages(messages)
                .user("zhaojinwei")
                .temperature(0.1d)
                .maxTokens(120)
                .stop(Arrays.asList("Human:", "AI:"))
                .n(1)
                .build();
        ChatCompletionResult chatCompletion = openAiService.createChatCompletion(chatCompletionRequest);
        List<ChatCompletionChoice> choices = chatCompletion.getChoices();
        for (ChatCompletionChoice choice : choices) {
            return choice.getMessage().getContent();
        }
        return "";
    }

    public String createImage(String text) {
        CreateImageRequest request = CreateImageRequest.builder()
                .prompt(text)
                .build();
        List<Image> data = openAiService.createImage(request).getData();
        return data.get(0).getUrl();
    }

    public String edit(String text, String editText) {
        EditRequest editRequest = EditRequest.builder()
                .model(editModel)
                .temperature(0.0)
                .input(text)
                .instruction(editText)
                .n(1)
                .build();
        EditResult edit = openAiService.createEdit(editRequest);
        List<EditChoice> choices = edit.getChoices();
        for (EditChoice choice : choices) {
            return choice.getText();
        }
        return text;
    }
}
